function [ y, alpha, beta, err ] = raykar( Y, Omega, X )
% Y:[L N] {0, 1}
% X:[D N]
% Omega: [L N], observed entries

N = size(Y, 2);

if(exist('X', 'var'))
    D = size(X, 1);
    w = randn(D, 1);
end

% init with major voting
mu = rand(N, 1);

maxIter = 100;
err = zeros(maxIter, 1);
for i = 1:maxIter
    %update w
    if(exist('X', 'var'))
        w  = updateW(mu, X, w);
    end
    
    % update alpha and beta
    alpha = updateAlpha(mu, Y, Omega);
    beta  = updateAlpha(1 - mu, 1 - Y, Omega);
    
    % update mu
    if(exist('X', 'var'))
        P = sigmoid(w'*X);
    else
        P = mean(mu);
    end
    mu_i = E_Step(Y, Omega, P, alpha, beta);
    
    err(i) = norm(mu_i - mu, 2);
    fprintf('iter:%d, err: %d\n', i, err(i));
    if(err(i) < 1e-6)
        break;
    end
    
    mu = mu_i;
end
err = err(1:i);

y = mu;

end

%% ---------------------------------------------------------------
function [w, err] = updateW(mu, X, initw)

maxIter = 1000;
w = initw;
for i = 1:maxIter
    sig = sigmoid(w'*X);
    grad = bsxfun(@times, X, mu' - sig);
    grad = sum(grad, 2);
    
    err = norm(grad, 2);
    % fprintf('iter:%d err:%d \n', i, err);
    if(err < 1e-4)
        break;
    end
    
    H = bsxfun(@times, sig, X);
    H = bsxfun(@times, 1- sig, X)*H';
    H = pinv(H); 
    
    w = w + 0.1*H*grad;
end

end

%% ---------------------------------------------------------------
function [mu] = E_Step(Y, Omega, P, alpha, beta)
N = size(Y, 2);
P = P';

A = zeros(N, 1);
B = zeros(N, 1);
for n = 1:N
    o = Omega(:, n);
    A(n) = getAlpha(alpha, Y(:, n), o);
    B(n) = getAlpha(beta, 1 - Y(:, n), o);
end

mu = (P.*A)./((P.*A) + B.*(1 - P));

end

%% ---------------------------------------------------------------
function [A] = updateAlpha(mu, Y, Omega)

L = size(Y, 1);
A = zeros(L, 1);
for l = 1:L
    y = Y(l, :);
    o = Omega(l, :);
    
    A(l) = sum(mu'.*y.*o)/sum(mu'.*o);
end

end

%% ---------------------------------------------------------------
function [val] = getAlpha(a, y, o)

a = a(o);
y = y(o);
val1 = a.^y;
val1 = prod(val1);

val0 = (1 - a).^(1 - y);
val0 = prod(val0);

val = val0*val1;

end

%% ---------------------------------------------------------------
function [a] = sigmoid(a)
a = 1./(1 + exp(-a));
end
